package cn.fly.chatgpt.config.server;

import cn.fly.chatgpt.config.WebSocketConfiguration;
import cn.fly.chatgpt.domain.entity.Chatgpt;
import cn.fly.chatgpt.domain.entity.Dialogue;
import cn.fly.chatgpt.domain.entity.Qus;
import cn.fly.chatgpt.service.ChatgptService;
import cn.hutool.json.JSON;
import cn.hutool.json.JSONUtil;
import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.databind.DeserializationFeature;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.PropertyNamingStrategy;
import com.ruoyi.common.core.utils.StringUtils;
import com.zhipu.oapi.ClientV4;
import com.zhipu.oapi.Constants;
import com.zhipu.oapi.service.v4.model.*;
import io.reactivex.Flowable;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import javax.websocket.*;
import javax.websocket.server.PathParam;
import javax.websocket.server.ServerEndpoint;
import java.io.IOException;
import java.time.Duration;
import java.time.LocalDateTime;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.regex.Matcher;
import java.util.regex.Pattern;


@Component
@ServerEndpoint(value = "/chatgpt/{sid}")
public class WebSocketServer {

    private static ChatgptService chatgptService;

    @Autowired
    public void setChatgptService(ChatgptService chatgptService) {
        WebSocketServer.chatgptService = chatgptService;
    }


    private static final String API_SECRET_KEY = "3f2e78893f9529a376112544d0244546.sEUa391TTH3HKy18";

    private static final ClientV4 client = new ClientV4.Builder(API_SECRET_KEY).build();

    // 自己的业务id
    private static final String requestIdTemplate = "Frankiefly";

    private static final ObjectMapper mapper = defaultObjectMapper();

    public static Flowable<ChatMessageAccumulator> mapStreamToAccumulator(Flowable<ModelData> flowable) {
        return flowable.map(chunk -> {
            return new ChatMessageAccumulator(chunk.getChoices().get(0).getDelta(), null, chunk.getChoices().get(0), chunk.getUsage(), chunk.getCreated(), chunk.getId());
        });
    }

    public static ObjectMapper defaultObjectMapper() {
        ObjectMapper mapper = new ObjectMapper();
        mapper.configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false);
        mapper.setSerializationInclusion(JsonInclude.Include.NON_NULL);
        mapper.setPropertyNamingStrategy(PropertyNamingStrategy.SNAKE_CASE);
        mapper.addMixIn(ChatFunction.class, ChatFunctionMixIn.class);
        mapper.addMixIn(ChatCompletionRequest.class, ChatCompletionRequestMixIn.class);
        mapper.addMixIn(ChatFunctionCall.class, ChatFunctionCallMixIn.class);
        return mapper;
    }

    private static StringRedisTemplate stringRedisTemplate;

    @Autowired
    public void setStringRedisTemplate(StringRedisTemplate stringRedisTemplate) {
        WebSocketServer.stringRedisTemplate = stringRedisTemplate;
    }

    private static Map<String, Session> sessionMap = new HashMap<>();

    @OnOpen
    public void onOpen(Session session, @PathParam("sid") String sid) {
        System.out.println("客户端: " + sid + "建立连接");
        sessionMap.put(sid, session);
    }

    @OnMessage
    public void onMessage(String askJson, @PathParam("sid") String sid){
        Integer userId = 0;
        Long chatId = 0L;

        Pattern pattern = Pattern.compile("u(\\d+)c(\\d+)");
        Matcher matcher = pattern.matcher(sid);

        if (matcher.find()) {
            String firstNumberStr = matcher.group(1); // 获取第一个数字的字符串形式
            String secondNumberStr = matcher.group(2); // 获取第二个数字的字符串形式

            userId = Integer.valueOf(firstNumberStr);
            chatId = Long.valueOf(secondNumberStr);

        } else
            System.out.println("出现错误");

        Qus qus = JSONUtil.parse(askJson).toBean(Qus.class);

        System.out.println(askJson);

        StringBuffer res = new StringBuffer();

        String question = qus.getQuestion();
        LocalDateTime qusTime = qus.getQusTime();
        System.out.println(question);

        List<ChatMessage> messages = new ArrayList<>();

        // 上下文
        String history = stringRedisTemplate.opsForValue().get("chatHistory" + chatId);
        ChatMessage chatHistory;
        if (StringUtils.hasText(history)) {
            chatHistory = new ChatMessage(ChatMessageRole.USER.value(), history);
            messages.add(chatHistory);
        }

        // 设置问题
        ChatMessage chatMessage = new ChatMessage(ChatMessageRole.USER.value(), question);

        messages.add(chatMessage);

        // 业务 id + 时间戳
        String requestId = String.format(requestIdTemplate, System.currentTimeMillis());

        // 函数调用参数构建部分
        List<ChatTool> chatToolList = new ArrayList<>();
        ChatTool chatTool = new ChatTool();
        chatTool.setType(ChatToolType.FUNCTION.value());

        ChatFunctionParameters chatFunctionParameters = new ChatFunctionParameters();
        chatFunctionParameters.setType("object");

        Map<String, Object> properties = new HashMap<>();
        properties.put("departure", new HashMap<String, Object>() {{
            put("type", "string");
            put("description", "description");
        }});
        properties.put("destination", new HashMap<String, Object>() {{
            put("type", "string");
            put("description", "description");
        }});
        properties.put("date", new HashMap<String, Object>() {{
            put("type", "string");
            put("description", "description");
        }});
        List<String> required = new ArrayList<>();
        required.add("departure");
        required.add("destination");
        required.add("date");
        chatFunctionParameters.setProperties(properties);
        ChatFunction chatFunction = ChatFunction.builder()
                .name("ChatGLM")
                .description("description")
                .parameters(chatFunctionParameters)
                .required(required)
                .build();
        chatTool.setFunction(chatFunction);
        chatToolList.add(chatTool);

        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .model(Constants.ModelChatGLM4)
                .stream(Boolean.TRUE)
                .messages(messages)
                .requestId(requestId)
                .tools(chatToolList)
                .toolChoice("auto")
                .build();
        ModelApiResponse sseModelApiResp = client.invokeModelApi(chatCompletionRequest);
        if (sseModelApiResp.isSuccess()) {
            AtomicBoolean isFirst = new AtomicBoolean(true);
            ChatMessageAccumulator chatMessageAccumulator = mapStreamToAccumulator(sseModelApiResp.getFlowable())
                    .doOnNext(accumulator -> {
                        {
                            if (isFirst.getAndSet(false)) {
                                System.out.print("Response: ");
                            }
                            if (accumulator.getDelta() != null && accumulator.getDelta().getTool_calls() != null) {
                                String jsonString = mapper.writeValueAsString(accumulator.getDelta().getTool_calls());
                                System.out.println("tool_calls: " + jsonString);
                            }
                            if (accumulator.getDelta() != null && accumulator.getDelta().getContent() != null) {
                                System.out.print(accumulator.getDelta().getContent());
                                sendToClient(sid, accumulator.getDelta().getContent());
                                res.append(accumulator.getDelta().getContent());
                            }
                        }
                    })
                    .doOnComplete(System.out::println)
                    .lastElement()
                    .blockingGet();

            LocalDateTime ansTime = LocalDateTime.now();

            stringRedisTemplate.opsForValue().set("chatHistory" + chatId, history + " " + question + " " + res, Duration.ofMinutes(10));

            Chatgpt chatgpt = new Chatgpt(question, qusTime, res.toString(), ansTime, chatId, userId);

            chatgptService.addInfo(chatgpt);

            sendToClient(sid, "ayiayibaba" + ansTime);

            Choice choice = new Choice(chatMessageAccumulator.getChoice().getFinishReason(), 0L, chatMessageAccumulator.getDelta());
            List<Choice> choices = new ArrayList<>();
            choices.add(choice);
            ModelData data = new ModelData();
            data.setChoices(choices);
            data.setUsage(chatMessageAccumulator.getUsage());
            data.setId(chatMessageAccumulator.getId());
            data.setCreated(chatMessageAccumulator.getCreated());
            data.setRequestId(chatCompletionRequest.getRequestId());
            sseModelApiResp.setFlowable(null);
            sseModelApiResp.setData(data);
        }


        //sendToClient(sid, "giao");
    }

    @OnClose
    public void onClose(@PathParam("sid") String sid){
        System.out.println("连接断开: " + sid);
        sessionMap.remove(sid);
    }

    public void sendToClient(String sid, String result){
        Session session = sessionMap.get(sid);
        try {
            session.getBasicRemote().sendText(result);
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }
}
